#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
豆包模型接口封装
负责与豆包AI模型API交互
"""

import os
import json
import yaml
import requests
import logging
from typing import Dict, Any, Optional

logger = logging.getLogger(__name__)

class APIClient:
    """
    API客户端，负责与豆包模型进行交互
    """
    
    def __init__(self, config_path: str = None):
        """
        初始化API客户端
        
        Args:
            config_path: 配置文件路径，默认为config/model_config.yaml
        """
        self.config_path = config_path or os.path.join("config", "model_config.yaml")
        self.config = self._load_config()
        self.api_key = self.config.get("api_key") or os.environ.get("DOUBAN_API_KEY")
        self.base_url = self.config.get("base_url", "https://api.doubao.com/v1")
        
        # 验证API密钥
        if not self.api_key:
            logger.warning("API密钥未配置，部分功能可能无法使用")
    
    def generate_answer(self, question: str, model_version: str = "v2") -> str:
        """
        生成回答文本
        
        Args:
            question: 问题文本
            model_version: 模型版本
            
        Returns:
            str: 生成的回答文本
            
        Raises:
            ValueError: 问题为空或API密钥未设置
            ConnectionError: API连接失败
            Exception: 其他API异常
        """
        # 检查问题是否为空
        if not question or question.strip() == "":
            raise ValueError("问题文本不能为空")
        
        # 检查API密钥
        if not self.api_key:
            raise ValueError("API密钥未设置，无法调用模型")
        
        # 获取模型URL
        model_config = self.config.get("models", {}).get(model_version, {})
        if not model_config:
            logger.warning(f"模型版本 {model_version} 配置不存在，使用默认配置")
            model_endpoint = "text/completions"
            model_params = {"temperature": 0.7, "max_tokens": 200}
        else:
            model_endpoint = model_config.get("endpoint", "text/completions")
            model_params = model_config.get("parameters", {})
        
        # 构建请求URL
        url = f"{self.base_url}/{model_endpoint}"
        
        # 构建请求体
        payload = {
            "model": f"doubao-{model_version}",
            "prompt": self._format_prompt(question),
            **model_params
        }
        
        # 构建请求头
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }
        
        try:
            # 发送请求
            logger.info(f"正在请求豆包模型API，版本: {model_version}")
            
            response = requests.post(
                url,
                headers=headers,
                json=payload,
                timeout=30  # 30秒超时
            )
            
            # 检查响应状态
            response.raise_for_status()
            
            # 解析响应
            result = response.json()
            
            # 提取生成的文本
            if "choices" in result and len(result["choices"]) > 0:
                answer = result["choices"][0]["text"].strip()
                logger.info(f"成功获取豆包模型回答 (长度: {len(answer)})")
                return answer
            else:
                logger.error(f"API响应格式异常: {result}")
                raise Exception("API响应格式不符合预期")
                
        except requests.exceptions.RequestException as e:
            logger.error(f"API请求失败: {e}")
            raise ConnectionError(f"无法连接到豆包模型API: {e}")
            
        except Exception as e:
            logger.error(f"生成回答时发生异常: {e}")
            raise Exception(f"调用豆包模型API时发生异常: {e}")
    
    def _load_config(self) -> Dict[str, Any]:
        """
        加载配置文件
        
        Returns:
            Dict: 配置信息
        """
        try:
            with open(self.config_path, 'r', encoding='utf-8') as f:
                config = yaml.safe_load(f)
                logger.info(f"成功加载豆包模型配置: {self.config_path}")
                return config
        except Exception as e:
            logger.error(f"加载配置文件失败: {e}, 使用默认配置")
            # 默认配置
            return {
                "base_url": "https://api.doubao.com/v1",
                "models": {
                    "v2": {
                        "endpoint": "text/completions",
                        "parameters": {
                            "temperature": 0.7,
                            "max_tokens": 200,
                            "top_p": 0.8
                        }
                    }
                }
            }
    
    def _format_prompt(self, question: str) -> str:
        """
        格式化提示文本
        
        Args:
            question: 原始问题
            
        Returns:
            str: 格式化后的提示文本
        """
        # 根据问题类型添加引导词
        if "评价" in question or "建议" in question:
            return (
                f"请针对以下问题提供详细、积极、客观的评价，字数在100-200字之间:\n\n"
                f"{question}\n\n"
                f"要求：\n"
                f"1. 态度积极，但要客观公正\n"
                f"2. 使用书面语，避免口语化表达\n"
                f"3. 不要使用过于夸张的赞美词汇\n"
                f"4. 如果有建议，应当委婉表达"
            )
        else:
            return question 